import torch
import torchvision
import random
import cv2
import numpy as np
import torch.nn as nn
import torchvision.transforms as transforms
import torch.nn.functional as F
from test import load_model
from numpy import linalg as LA
import math
import time
import sys
from PIL import Image
import torch.utils.data as data
from matplotlib import pyplot as plt
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
loss_func = nn.CrossEntropyLoss()


def show_img(img):
    img = img[0]
    img = img.cpu().detach().numpy()
    img = img.transpose(1, 2, 0)
    plt.imshow(img)
    plt.show()


class OriDataset(torch.utils.data.Dataset):
    def __init__(self, transform):
        images = np.load('./data_select.npy')
        labels = np.load('./label_select.npy')
        assert labels.min() >= 0
        assert images.dtype == np.uint8
        self.images = [Image.fromarray(x) for x in images]
        self.labels = labels / labels.sum(axis=1, keepdims=True)  # normalize
        self.labels = self.labels.astype(np.float32)
        self.transform = transform

    def __getitem__(self, index):
        image, label = self.images[index], self.labels[index]
        image = self.transform(image)
        return image, label

    def __len__(self):
        return len(self.labels)


# PGD攻击
def pgd_attack(models, image, label, eps=8 / 255, alpha=1 / 255, iters=10):

    adv_image = torch.clone(image).to(device)
    # 噪声初始化, TODO 均匀噪声、低频噪声
    noise = np.random.uniform(-eps, eps, adv_image.shape)
    adv_image = adv_image + torch.tensor(noise, device=device)
    adv_image = torch.clamp(adv_image, 0, 1.0)

    for i in range(iters):
        cal_img = normalize(adv_image)
        cal_img = cal_img.type(torch.cuda.FloatTensor)
        # cal_img = cal_img.float()
        cal_img.requires_grad = True

        cost = torch.tensor(0, device=device, dtype=torch.double)
        for model in models:
            outputs = model(cal_img)
            model.zero_grad()
            targets = label.argmax(dim=-1)
            # targets = targets.unsqueeze(0)
            cost += loss_func(outputs, targets)

        # cost = cost.double()
        cost = cost.double()
        cost.backward()

        adv_image = adv_image + alpha * cal_img.grad.sign()
        eta = torch.clamp(adv_image - image, min=-eps, max=eps)
        adv_image = torch.clamp(image + eta, min=0, max=1).detach_()

    return adv_image


# 随机旋转图片
def roation(image):
    ori = torch.clone(image)
    angle = random.randint(0, 180) * math.pi / 180
    theta = torch.tensor([
        [math.cos(angle), math.sin(-angle), 0],
        [math.sin(angle), math.cos(angle), 0]
    ], dtype=torch.double, device=device)
    grid = F.affine_grid(theta.unsqueeze(0), image.size())
    output = F.grid_sample(image, grid)
    return ori, output


# 对图像做shear变换
def shear(image):
    ori = torch.clone(image)
    theta = torch.tensor([
        [1, random.random(), 0],
        [0, 1, 0]
    ], dtype=torch.double, device=device)
    grid = F.affine_grid(theta.unsqueeze(0), image.size())
    output = F.grid_sample(image, grid)
    return ori, output


# 对图像做translate变换
def translate(image):
    ori = torch.clone(image)
    theta = torch.tensor([
        [1, 0, random.uniform(0, 0.3)],
        [0, 1, random.uniform(0, 0.3)]
    ], dtype=torch.double, device=device)
    grid = F.affine_grid(theta.unsqueeze(0), image.size())
    output = F.grid_sample(image, grid)
    return ori, output


def affine_transform(img):
    img0, img1 = roation(img)

    img00, img01 = shear(img0)
    img10, img11 = shear(img1)

    img000, img001 = translate(img00)
    img010, img011 = translate(img01)
    img100, img101 = translate(img10)
    img110, img111 = translate(img11)

    return img000, img001, img010, img011, img100, img101, img110, img111


# 对图像做分块、置乱
def RFL(img):
    """
    将图像划分为16个8*8的图像块
    将16个图像块打乱顺序
    """
    block_list = []
    for w in range(4):
        for h in range(4):
            block_list.append(img[0, :, w*8:(w+1)*8, h*8:(h+1)*8])
    random.shuffle(block_list)
    shuffle_img = torch.zeros(size=img.shape, device=device)
    for w in range(4):
        for h in range(4):
            shuffle_img[0, :, w*8:(w+1)*8, h*8:(h+1)*8] = block_list[w*4+h]
    return shuffle_img


def change_and_append(img, enhance_images):
    img = img[0]
    img = img * 255
    img = img.cpu().detach().numpy()
    img = img.transpose(1, 2, 0)
    enhance_images.append(img)
    return enhance_images


if __name__ == '__main__':
    print('start!')
    start = time.time()

    models = []

    resnet50 = load_model('F:\\AliTianChi\\ali_8_attack\\resnet50.pth.tar', ResTag=True)
    resnet50.eval()
    models.append(resnet50)

    densenet121 = load_model('F:\\AliTianChi\\ali_8_attack\\densenet121.pth.tar', ResTag=False)
    densenet121.eval()
    models.append(densenet121)

    transform_train = transforms.Compose([
        transforms.ToTensor(),
    ])
    trainset = OriDataset(transform=transform_train)
    trainloader = data.DataLoader(trainset, batch_size=1, shuffle=False, num_workers=0)

    enhance_imgs = []
    enhance_labels = []
    count = 0
    for (inputs, soft_labels) in trainloader:
        count += 1
        if count % 10 == 0:
            print(f'count={count} cost_time={time.time()-start}s')
        inputs = inputs.to(device)
        soft_labels = soft_labels.to(device)

        # show_img(img=inputs)
        ori = torch.clone(inputs)
        enhance_imgs = change_and_append(img=ori, enhance_images=enhance_imgs)
        # 先对图像进行PGD攻击
        advs = pgd_attack(models=models, image=inputs, label=soft_labels)
        # show_img(img=advs)
        # 使用对抗样本 进行 旋转、希尔、平移 ，得到组合的8种方式
        img000, img001, img010, img011, img100, img101, img110, img111 = affine_transform(img=advs)
        enhance_imgs = change_and_append(img=img000, enhance_images=enhance_imgs)
        enhance_imgs = change_and_append(img=img001, enhance_images=enhance_imgs)
        enhance_imgs = change_and_append(img=img010, enhance_images=enhance_imgs)
        enhance_imgs = change_and_append(img=img011, enhance_images=enhance_imgs)
        enhance_imgs = change_and_append(img=img100, enhance_images=enhance_imgs)
        enhance_imgs = change_and_append(img=img101, enhance_images=enhance_imgs)
        enhance_imgs = change_and_append(img=img110, enhance_images=enhance_imgs)
        enhance_imgs = change_and_append(img=img111, enhance_images=enhance_imgs)
        # show_img(img=img000)
        # show_img(img=img001)
        # show_img(img=img010)
        # show_img(img=img011)
        # show_img(img=img100)
        # show_img(img=img101)
        # show_img(img=img110)
        # show_img(img=img111)
        # 使用对抗样本 进行 RLF变换
        shuffle_img = RFL(img=advs)
        enhance_imgs = change_and_append(img=shuffle_img, enhance_images=enhance_imgs)
        # show_img(img=shuffle_img)

        # 执行10次
        labels = soft_labels.cpu().numpy()[0]
        for c in range(10):
            enhance_labels.append(labels)

    data_select = np.asarray(enhance_imgs, dtype=np.uint8)
    label_select = np.asarray(enhance_labels)

    print(data_select.shape)
    print(label_select.shape)
    np.save('data.npy', data_select)
    np.save('label.npy', label_select)
    print(f'cost_time={time.time()-start}s')
    exit()

